from transformers import BertTokenizer

#Load the fine tuned model from the path
PATH = 'test_trainer/'
md = AutoModelForSequenceClassification.from_pretrained(PATH, local_files_only=True)

def make_classification(text):
    # Tokenize
    inps = brt_tkn(text, padding=True, truncation=True, max_length=512, return_tensors="pt").to("cuda")
    # get output
    outputs = model(**inps)
    # softmax for generating probablities
    probablities = outputs[0].softmax(1)
    # get best match.
    return probablities .argmax()

#First inference
text = """
This is the show that puts a smile on your face as you watch it. You get in love with each and every character of the show. At the end, I felt eight episode were not enough. Will wait for season 2.
"""
print(make_classification(text))
